import os
import argparse
import time
import pickle
import shutil

import numpy as np
import tensorflow as tf
import keras.backend as K
from keras.optimizers import SGD, RMSprop, Adam

from VGG13_v9 import *


def build_model():
    sub_model = vgg13_shortcuts_v2(vgg13_config.input_shape, vgg13_config.base_nb_filters, vgg13_config.weight_decay, init='he_normal', drop_rate=0.5)
    main_model = Model(input=sub_model.inputs[0], output=[sub_model.outputs[0], sub_model.get_layer(name='feature').output])
    return main_model, K.function([sub_model.inputs[0], K.learning_phase()], [sub_model.get_layer(name='feature').output])


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="vgg13_v9 with center loss")

    parser.add_argument("-r", "--lr",
                        type=float, dest="lr", metavar="lr",
                        default=vgg13_config.lr, required=False,
                        help="learning rate")

    parser.add_argument("-g", "--gpu", required=True,
                        dest="gpu", metavar="gpu",
                        help="gpu device to train")

    parser.add_argument("-t", "--test", required=True,
                        type=int, dest="test_flag", metavar="test flag",
                        help="flag whether to test or validate")

    parser.add_argument("-f", "--fold", required=True,
                        type=int, dest="valid_fold", metavar="valid fold index",
                        help="fold index of model")

    parser.add_argument("-i", "--include", required=True,
                        type=int, dest="include_test", metavar="include test flag",
                        help="flag whether to include kaggle testset")

    parser.add_argument("-o", "--output", required=True,
                        dest="output_dir", metavar="output dir",
                        help="directory to output result")

    parser.add_argument("-w", "--weights", required=True,
                        dest="weights_path", metavar="model weights",
                        help="model weights path")

    parser.add_argument("-d", "--hdf5", required=True,
                        dest="hdf5_path", metavar="candidates volume hdf5",
                        help="candidates volume path")

    parser.add_argument("-p", "--pkl", required=True,
                        dest="candidates_path", metavar="candidates info path",
                        help="candidates info pkl path")

    args = parser.parse_args()

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    tf_config = tf.ConfigProto()
    tf_config.log_device_placement = False
    tf_config.gpu_options.allow_growth = True
    K.set_session(tf.Session(config=tf_config))

    # build the model
    model = vgg13_shortcuts_v2(vgg13_config.input_shape, vgg13_config.base_nb_filters, vgg13_config.weight_decay,
                               init='he_normal', drop_rate=0)

    assert args.weights_path is not None
    weights_path = os.listdir(args.weights_path)
    assert len(weights_path) > 0, 'require at least one model in the directory'
    weights_path = os.path.join(args.weights_path, weights_path[0])
    model.load_weights(weights_path, by_name=True)
    print('successfully load weights from %s' % weights_path)

    # model.summary()

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    if args.test_flag:
        gen = Test_Generator(args.hdf5_path, args.candidates_path)
    else:
        gen = Valid_Generator(args.hdf5_path, args.candidates_path, args.valid_fold, args.include_test)

    # not reverse
    val_generator = gen.val_generate(8)
    print('predict begins.')
    preds = model.predict_generator(val_generator, len(gen.val_candidates), max_q_size=8)
    val_generator.close()

    # # reverse
    # val_generator = gen.val_generate(8, reverse=True)
    # print('predict(reverse) begins.')
    # preds2 = model.predict_generator(val_generator, len(gen.val_candidates), max_q_size=8)
    # val_generator.close()

    # for ind, candidate in enumerate(gen.val_candidates):
    #     candidate['mal_prob'] = max(preds[ind, 0], preds2[ind, 0])
    for ind, candidate in enumerate(gen.val_candidates):
        candidate['mal_prob'] = preds[ind, 0]

    if args.test_flag:
        with open(os.path.join(args.output_dir, 'vgg13_v9_merge%d.pkl' % args.valid_fold), 'wb') as f:
            pickle.dump(gen.val_candidates, f, -1)
    else:
        with open(os.path.join(args.output_dir, 'vgg13_v9_cv_fold%d.pkl' % args.valid_fold), 'wb') as f:
            pickle.dump(gen.val_candidates, f, -1)

    K.clear_session()
    print('predict ends.')